import numpy as np
from Xlib import display, X
import cv2
import time
import socket
import pickle
import struct

class ScreenServer:
    def __init__(self, client_ip="192.168.10.26", port=5005, jpeg_quality=60, max_packet_size=1400):
        # UDP配置
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.client_addr = (client_ip, port)
        
        # 屏幕参数
        self.disp = display.Display()
        self.screen = self.disp.screen()
        self.root = self.screen.root
        self.width = self.screen.width_in_pixels
        self.height = self.screen.height_in_pixels
        self.format = X.ZPixmap
        
        # 块分割（4x4=16块）
        self.grid_rows = 4
        self.grid_cols = 4
        self.block_width = self.width // self.grid_cols
        self.block_height = self.height // self.grid_rows
        
        # 压缩和分片参数
        self.jpeg_quality = jpeg_quality
        self.encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), self.jpeg_quality]
        self.max_packet_size = max_packet_size
        
        # 内存存储最新块（不写入磁盘）
        self.latest_blocks = {}  # (i,j) -> 原始块图像
        self.first_frame_sent = False
        self.packet_id = 0  # 数据包ID

    def get_mouse_position(self):
        """获取鼠标坐标"""
        data = self.root.query_pointer()
        return (data.root_x, data.root_y)

    def capture_frame(self):
        """捕获屏幕帧（不含鼠标）"""
        try:
            raw = self.root.get_image(
                x=0, y=0,
                width=self.width, height=self.height,
                format=self.format,
                plane_mask=0xFFFFFFFF
            )
            
            # 处理行字节对齐
            try:
                actual_row_bytes = raw.bytes_per_line
            except AttributeError:
                actual_row_bytes = len(raw.data) // self.height

            # 转换为BGR格式
            data = np.frombuffer(raw.data, dtype=np.uint8)
            frame = []
            for y in range(self.height):
                start = y * actual_row_bytes
                end = start + self.width * 4  # 32位格式，4字节/像素
                row = data[start:end]
                if len(row) < self.width * 4:
                    row = np.pad(row, (0, self.width * 4 - len(row)), mode='constant')
                frame.append(row)
            
            frame = np.concatenate(frame).reshape((self.height, self.width, 4))
            return cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
        except Exception as e:
            print(f"捕获帧失败: {e}")
            return None

    def split_into_blocks(self, frame):
        """分割帧为16块"""
        blocks = []
        for i in range(self.grid_rows):
            for j in range(self.grid_cols):
                y1 = i * self.block_height
                y2 = (i + 1) * self.block_height if i < self.grid_rows - 1 else self.height
                x1 = j * self.block_width
                x2 = (j + 1) * self.block_width if j < self.grid_cols - 1 else self.width
                blocks.append((i, j, frame[y1:y2, x1:x2]))
        return blocks

    def compress_block(self, block):
        """压缩块为JPEG字节流（仅内存操作）"""
        result, encoded = cv2.imencode('.jpg', block, self.encode_param)
        return encoded.tobytes() if result else None

    def get_changed_blocks(self, current_frame):
        """检测变化的块（内存中比较）"""
        current_blocks = self.split_into_blocks(current_frame)
        changed = []
        
        for i, j, curr_block in current_blocks:
            if (i, j) not in self.latest_blocks:
                # 新块（第一帧）
                self.latest_blocks[(i, j)] = curr_block.copy()
                compressed = self.compress_block(curr_block)
                if compressed:
                    changed.append((i, j, compressed))
            else:
                # 比较差异
                prev_block = self.latest_blocks[(i, j)]
                if np.mean(cv2.absdiff(curr_block, prev_block)) > 5.0:
                    self.latest_blocks[(i, j)] = curr_block.copy()
                    compressed = self.compress_block(curr_block)
                    if compressed:
                        changed.append((i, j, compressed))
        
        return changed

    def split_into_packets(self, data):
        """数据分片"""
        packets = []
        data_len = len(data)
        offset = 0
        
        while offset < data_len:
            chunk_size = min(self.max_packet_size, data_len - offset)
            chunk = data[offset:offset + chunk_size]
            total_packets = (data_len + self.max_packet_size - 1) // self.max_packet_size
            header = struct.pack(">IIII", self.packet_id, total_packets, len(packets), chunk_size)
            packets.append(header + chunk)
            offset += chunk_size
        
        self.packet_id = (self.packet_id + 1) % 65536
        return packets

    def send_data(self, mouse_pos, changed_blocks=None):
        """发送数据（仅内存数据）"""
        data = {"mouse": mouse_pos}
        
        if not self.first_frame_sent:
            data["screen_info"] = (self.width, self.height)
            data["changed_blocks"] = changed_blocks
            self.first_frame_sent = True
        else:
            data["changed_blocks"] = changed_blocks if changed_blocks else []
        
        try:
            serialized = pickle.dumps(data)
            packets = self.split_into_packets(serialized)
            for packet in packets:
                self.udp_socket.sendto(packet, self.client_addr)
            return True
        except Exception as e:
            print(f"发送失败: {e}")
            return False

    def run(self):
        print(f"服务器启动，目标: {self.client_addr} | 分辨率: {self.width}x{self.height}")
        try:
            while True:
                frame = self.capture_frame()
                if frame is None:
                    time.sleep(0.1)
                    continue
                
                changed_blocks = self.get_changed_blocks(frame)
                mouse_pos = self.get_mouse_position()
                self.send_data(mouse_pos, changed_blocks)
                
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("服务器停止")
        finally:
            self.udp_socket.close()

if __name__ == "__main__":
    server = ScreenServer(client_ip="192.168.10.26")
    server.run()